import torch
from torch import nn


class StockLSTM(nn.Module):

    def __init__(self, hidden_size, num_layers, output_size, input_dim):
        super().__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.input_dim = input_dim

        self.lstm = nn.LSTM(input_size=input_dim, hidden_size=hidden_size, num_layers=num_layers, batch_first=True)
        self.fc = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(x.device).requires_grad_()
        c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(x.device).requires_grad_()
        r_out, (h_n, h_c) = self.lstm(x, (h0, c0))
        out = self.fc(r_out[:, -1, :])
        return out
